package com.wangliang.deepseek.service;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.wangliang.deepseek.config.ApiConfig;
import com.wangliang.deepseek.constants.ApiConstants;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okio.Buffer;
import okio.BufferedSource;
import okio.ByteString;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * DeepSeek API 服务类
 * 用于处理与 DeepSeek API 的通信
 */

@Slf4j
@Service
public class DeepSeekService {

    private final OkHttpClient client;
    private final Gson gson;
    private final ExecutorService executorService;
    private final ApiConfig apiConfig;

    public DeepSeekService(ApiConfig apiConfig) {
        this.apiConfig = apiConfig;
        this.gson = new Gson();
        this.executorService = Executors.newCachedThreadPool();
        this.client = new OkHttpClient.Builder()
                .connectTimeout(apiConfig.getTimeoutSeconds(), TimeUnit.SECONDS)
                .readTimeout(apiConfig.getTimeoutSeconds(), TimeUnit.SECONDS)
                .writeTimeout(apiConfig.getTimeoutSeconds(), TimeUnit.SECONDS)
                .retryOnConnectionFailure(true)
                .addInterceptor(chain -> {
                    Request original = chain.request();
                    Request request = original.newBuilder()
                            .header("Content-Type", "application/json; charset=utf-8")
                            .header("Accept", "text/event-stream")
                            .header("Authorization", "Bearer " + apiConfig.getApiKey())
                            .method(original.method(), original.body())
                            .build();
                    return chain.proceed(request);
                })
                .build();
    }

    /**
     * 发送流式查询
     */
    public void sendStreamQuery(String query, Consumer<String> onMessage, Consumer<Throwable> onError) {
        if (StringUtils.isBlank(query)) {
            onError.accept(new IllegalArgumentException("Query cannot be empty"));
            return;
        }

        JsonObject requestBody = buildRequestBody(query, true);
        Request request = buildRequest(requestBody);

        // 使用线程池执行请求
        executorService.submit(() -> {
            int retries = 0;
            while (retries < apiConfig.getMaxRetries()) {
                try {
                    executeStreamRequest(request, onMessage);
                    return;
                } catch (Exception e) {
                    retries++;
                    if (retries >= apiConfig.getMaxRetries()) {
                        onError.accept(new IOException("Failed after " + apiConfig.getMaxRetries() + " retries", e));
                        return;
                    }
                    log.warn("Retry {}/{} after error: {}", retries, apiConfig.getMaxRetries(), e.getMessage());
                    try {
                        Thread.sleep(apiConfig.getRetryDelayMs() * retries);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        onError.accept(ie);
                        return;
                    }
                }
            }
        });
    }

    private void executeStreamRequest(Request request, Consumer<String> onMessage) throws IOException {
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected response: " + response);
            }

            ResponseBody body = response.body();
            if (body == null) {
                throw new IOException("Response body is null");
            }

            BufferedReader reader = new BufferedReader(
                    new InputStreamReader(body.byteStream(), StandardCharsets.UTF_8));
            String line;
            while ((line = reader.readLine()) != null) {
                if (line.startsWith("data: ")) {
                    String data = line.substring(6).trim();
                    if ("[DONE]".equals(data)) {
                        break;
                    }
                    try {
                        JsonObject jsonResponse = gson.fromJson(data, JsonObject.class);
                        String content = extractContent(jsonResponse);
                        if (content != null) {
                            onMessage.accept(content);
                        }
                    } catch (Exception e) {
                        log.warn("Failed to parse SSE message: {}", data, e);
                    }
                }
            }
        }
    }

    private JsonObject buildRequestBody(String query, boolean stream) {
        JsonObject message = new JsonObject();
        message.addProperty("role", "user");
        message.addProperty("content", query);

        JsonObject requestBody = new JsonObject();
        requestBody.addProperty("model", ApiConstants.MODEL_NAME);
        requestBody.add("messages", gson.toJsonTree(new JsonObject[] { message }));
        requestBody.addProperty("stream", stream);
        requestBody.addProperty("temperature", 0.7);
        requestBody.addProperty("max_tokens", 2000);

        return requestBody;
    }

    private Request buildRequest(JsonObject requestBody) {
        return new Request.Builder()
                .url(ApiConstants.BASE_URL)
                .post(RequestBody.create(requestBody.toString(), ApiConstants.JSON))
                .build();
    }

    private String extractContent(JsonObject jsonResponse) {
        try {
            return jsonResponse.getAsJsonArray("choices")
                    .get(0).getAsJsonObject()
                    .getAsJsonObject("delta")
                    .get("content").getAsString();
        } catch (Exception e) {
            log.debug("No content in delta: {}", jsonResponse);
            return null;
        }
    }

    /**
     * 发送查询到 DeepSeek API
     *
     * @param query 用户查询内容
     * @return API 响应的内容
     * @throws IllegalArgumentException 当查询参数为空时
     * @throws IOException              当 API 调用失败时
     */
    public String sendQuery(String query) throws IOException {
        // 参数校验
        if (StringUtils.isBlank(query)) {
            throw new IllegalArgumentException("Query cannot be empty");
        }

        // 构建请求体
        JsonObject requestBody = buildRequestBody(query, false);

        // 发送请求
        Request request = buildRequest(requestBody);

        return executeRequest(request);
    }

    /**
     * 关闭服务，释放资源
     */
    public void shutdown() {
        log.info("Shutting down DeepSeek service...");
        try {
            // 关闭线程池
            executorService.shutdown();
            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }

            // 关闭 OkHttpClient
            client.dispatcher().executorService().shutdown();
            client.connectionPool().evictAll();
            if (client.cache() != null) {
                client.cache().close();
            }

            log.info("DeepSeek service shutdown completed");
        } catch (Exception e) {
            log.error("Error during service shutdown", e);
            // 强制关闭线程池
            executorService.shutdownNow();
        }
    }

    private String executeRequest(Request request) throws IOException {
        int retries = 0;
        IOException lastException = null;

        while (retries < apiConfig.getMaxRetries()) {
            try (Response response = client.newCall(request).execute()) {
                ResponseBody body = response.body();
                if (body == null) {
                    throw new IOException("Response body is null");
                }

                if (!response.isSuccessful()) {
                    handleErrorResponse(response);
                }

                return extractResponseContent(response);
            } catch (IOException e) {
                lastException = e;
                retries++;
                if (retries < apiConfig.getMaxRetries()) {
                    log.warn("Request failed, retrying ({}/{}): {}",
                            retries, apiConfig.getMaxRetries(), e.getMessage());
                    try {
                        Thread.sleep(apiConfig.getRetryDelayMs());
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        throw new IOException("Request interrupted", ie);
                    }
                }
            }
        }

        throw new IOException("Request failed after " + apiConfig.getMaxRetries() + " retries", lastException);
    }

    private void handleErrorResponse(Response response) throws IOException {
        ResponseBody body = response.body();
        if (body == null) {
            throw new IOException("Error response body is null");
        }

        String errorBody = new String(body.bytes(), StandardCharsets.UTF_8);
        log.error("API error response: {} - {}", response.code(), errorBody);
        throw new IOException(String.format("API request failed with code %d: %s",
                response.code(), errorBody));
    }

    private String extractResponseContent(Response response) throws IOException {
        ResponseBody body = response.body();
        if (body == null) {
            throw new IOException("Response body is null");
        }

        try {
            String responseBody = new String(body.bytes(), StandardCharsets.UTF_8);
            JsonObject jsonResponse = gson.fromJson(responseBody, JsonObject.class);

            // 检查 JSON 响应的结构
            if (!jsonResponse.has("choices") ||
                    jsonResponse.getAsJsonArray("choices").size() == 0) {
                throw new IOException("Invalid API response format: " + responseBody);
            }

            return jsonResponse.getAsJsonArray("choices")
                    .get(0).getAsJsonObject()
                    .getAsJsonObject("message")
                    .get("content").getAsString();
        } catch (Exception e) {
            log.error("Failed to parse API response", e);
            throw new IOException("Failed to parse API response", e);
        }
    }

    public String formatAIResponse(String response) {
        StringBuilder formattedResponse = new StringBuilder();

        // 添加响应容器
        formattedResponse.append("<div class='ai-response-container'>");

        // 添加响应头部，包含关闭按钮和工具栏
        formattedResponse.append("<div class='ai-response-header'>");
        formattedResponse.append("<div class='ai-response-toolbar'>");
        formattedResponse.append("<button class='toolbar-button copy-button' onclick='copyContent(this)'>");
        formattedResponse.append("<i class='fas fa-copy'></i> 复制");
        formattedResponse.append("</button>");
        formattedResponse.append("<button class='toolbar-button close-button' onclick='closeAIResponse(this)'>");
        formattedResponse.append("<i class='fas fa-times'></i>");
        formattedResponse.append("</button>");
        formattedResponse.append("</div>"); // end toolbar
        formattedResponse.append("</div>"); // end header

        // 添加响应内容区域
        formattedResponse.append("<div class='ai-response-content'>");

        // 检查是否包含代码块，如果包含则添加特殊格式
        if (response.contains("```")) {
            formattedResponse.append(formatCodeBlocks(response));
        } else {
            formattedResponse.append("<div class='text-content'>").append(response).append("</div>");
        }

        formattedResponse.append("</div>"); // end content
        formattedResponse.append("</div>"); // end container

        return formattedResponse.toString();
    }

    /**
     * 格式化代码块
     */
    private String formatCodeBlocks(String response) {
        StringBuilder formatted = new StringBuilder();
        String[] parts = response.split("((?=```)|(?<=```))");
        boolean inCodeBlock = false;
        String currentLanguage = "";

        for (String part : parts) {
            if (part.startsWith("```")) {
                if (!inCodeBlock) {
                    // 开始代码块
                    currentLanguage = part.length() > 3 ? part.substring(3).trim() : "";
                    formatted.append("<div class='code-block-container'>");
                    if (!currentLanguage.isEmpty()) {
                        formatted.append("<div class='code-block-header'>");
                        formatted.append("<span class='code-language'>").append(currentLanguage).append("</span>");
                        formatted.append("</div>");
                    }
                    formatted.append("<pre class='code-block'><code class='language-").append(currentLanguage)
                            .append("'>");
                    inCodeBlock = true;
                } else {
                    // 结束代码块
                    formatted.append("</code></pre></div>");
                    inCodeBlock = false;
                }
            } else {
                if (inCodeBlock) {
                    // 在代码块内，进行HTML转义并保持格式
                    formatted.append(escapeHtml(part.trim()));
                } else {
                    // 在代码块外，作为普通文本处理
                    formatted.append("<p>").append(part.trim()).append("</p>");
                }
            }
        }

        return formatted.toString();
    }

    /**
     * HTML字符转义
     */
    private String escapeHtml(String text) {
        return text.replace("&", "&amp;")
                .replace("<", "&lt;")
                .replace(">", "&gt;")
                .replace("\"", "&quot;")
                .replace("'", "&#39;");
    }
}